# -*- coding: utf-8 -*-
# time: 2025/4/8 16:36
# file: demo.py
# author: hanson
import torch
"""
如果显存不足，尝试减小 max_length 或启用梯度累积
"""
from transformers import (
    AutoTokenizer,
    AutoModelForCausalLM,
    TrainingArguments,
    Trainer,
    DataCollatorForTokenClassification
)
from datasets import load_dataset
from peft import LoraConfig, get_peft_model, TaskType
import numpy as np
from sklearn.metrics import classification_report

# 1. 加载数据集（假设是类似CLUENER的格式）
dataset = load_dataset("clue", "CLUENER")  # 替换为你的数据集名称
label_list = dataset["train"].features["ner_tags"].feature.names
label2id = {label: i for i, label in enumerate(label_list)}
id2label = {i: label for i, label in enumerate(label_list)}

# 2. 加载Qwen-0.5的tokenizer和模型
model_name = "Qwen/Qwen-0.5B"  # 替换为实际模型路径
tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
tokenizer.pad_token = tokenizer.eos_token  # 设置pad_token

model = AutoModelForCausalLM.from_pretrained(model_name, trust_remote_code=True)
model.resize_token_embeddings(len(tokenizer))

# 3. 数据预处理
def tokenize_and_align_labels(examples):
    tokenized_inputs = tokenizer(
        examples["tokens"],
        truncation=True,
        is_split_into_words=True,
        padding="max_length",
        max_length=128
    )
    labels = []
    for i, label in enumerate(examples["ner_tags"]):
        word_ids = tokenized_inputs.word_ids(batch_index=i)
        previous_word_idx = None
        label_ids = []
        for word_idx in word_ids:
            if word_idx is None:
                label_ids.append(-100)
            elif word_idx != previous_word_idx:
                label_ids.append(label[word_idx])
            else:
                label_ids.append(-100)
            previous_word_idx = word_idx
        labels.append(label_ids)
    tokenized_inputs["labels"] = labels
    return tokenized_inputs

tokenized_datasets = dataset.map(
    tokenize_and_align_labels,
    batched=True,
    remove_columns=dataset["train"].column_names
)

# 4. 配置LoRA
peft_config = LoraConfig(
    task_type=TaskType.TOKEN_CLS,
    inference_mode=False,
    r=8,
    lora_alpha=32,
    lora_dropout=0.1,
    target_modules=["q_proj", "v_proj"]  # 针对Qwen的注意力层
)
model = get_peft_model(model, peft_config)
model.print_trainable_parameters()

# 5. 训练参数
training_args = TrainingArguments(
    output_dir="./qwen-ner-finetuned",
    per_device_train_batch_size=8,
    per_device_eval_batch_size=8,
    num_train_epochs=3,
    logging_dir="./logs",
    logging_steps=50,
    save_steps=500,
    evaluation_strategy="steps",
    eval_steps=500,
    learning_rate=5e-5,
    fp16=True,
    load_best_model_at_end=True,
)

# 6. 定义评估指标
def compute_metrics(p):
    predictions, labels = p
    predictions = np.argmax(predictions, axis=2)
    true_labels = [[id2label[l] for l in label if l != -100] for label in labels]
    true_predictions = [
        [id2label[p] for (p, l) in zip(prediction, label) if l != -100]
        for prediction, label in zip(predictions, labels)
    ]
    return classification_report(
        np.concatenate(true_labels),
        np.concatenate(true_predictions),
        output_dict=True
    )

# 7. 训练
trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_datasets["train"],
    eval_dataset=tokenized_datasets["validation"],
    data_collator=DataCollatorForTokenClassification(tokenizer),
    compute_metrics=compute_metrics
)
trainer.train()

# 8. 保存模型
model.save_pretrained("./qwen-ner-finetuned-final")